package com.jwong.auth.server.config;

import com.jwong.auth.server.model.UserModel;
import com.jwong.auth.server.sms.OverrideTokenGranterConfigurer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.config.annotation.configurers.ClientDetailsServiceConfigurer;
import org.springframework.security.oauth2.config.annotation.web.configuration.AuthorizationServerConfigurerAdapter;
import org.springframework.security.oauth2.config.annotation.web.configuration.EnableAuthorizationServer;
import org.springframework.security.oauth2.config.annotation.web.configurers.AuthorizationServerEndpointsConfigurer;
import org.springframework.security.oauth2.config.annotation.web.configurers.AuthorizationServerSecurityConfigurer;
import org.springframework.security.oauth2.provider.ClientDetailsService;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.client.JdbcClientDetailsService;
import org.springframework.security.oauth2.provider.code.AuthorizationCodeServices;
import org.springframework.security.oauth2.provider.code.JdbcAuthorizationCodeServices;
import org.springframework.security.oauth2.provider.token.AuthorizationServerTokenServices;
import org.springframework.security.oauth2.provider.token.DefaultTokenServices;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.store.JwtTokenStore;
import org.springframework.security.oauth2.provider.token.store.KeyStoreKeyFactory;
import org.springframework.web.bind.annotation.SessionAttributes;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;
import org.springframework.web.servlet.view.RedirectView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import javax.sql.DataSource;
import java.security.KeyPair;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

@Configuration
@SessionAttributes("authorizationRequest")
@EnableAuthorizationServer
public class AuthorizationServerConfiguration extends AuthorizationServerConfigurerAdapter {

    @Autowired
    private AuthenticationManager authenticationManager;

    @Autowired
    private PasswordEncoder passwordEncoder;

    @Autowired
    @Qualifier("clientDetailsDataSource")
    private DataSource clientDetailsDataSource;

    private static final String USER_PRIMARY_KEY = "user_primary_key";
    private static final String SCHOOL_CODE_KEY = "school_code_key";

    @Bean
    public JwtTokenStore tokenStore() {
        return new JwtTokenStore(accessTokenConverter());
    }

    @Bean
    public JwtAccessTokenConverter accessTokenConverter() {

        JwtAccessTokenConverter converter = new JwtAccessTokenConverter() {
            @Override
            public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
                // 将userID附加到accessToken
                UserModel principal = (UserModel)authentication.getUserAuthentication().getPrincipal();
                Optional.of(principal).ifPresent((user) -> {
                    final Map<String, Object> attached = new HashMap<>();
                    attached.put(USER_PRIMARY_KEY, user.getUserId());
                    Optional.ofNullable(user.getSchoolCode()).ifPresent(schoolCode -> attached.put(SCHOOL_CODE_KEY, schoolCode));
                    Map<String, Object> unmodifiableAttached = Collections.unmodifiableMap(attached);
                    ((DefaultOAuth2AccessToken)accessToken).setAdditionalInformation(unmodifiableAttached);
                });
                return super.enhance(accessToken, authentication);
            }
        };

        KeyPair keyPair = new KeyStoreKeyFactory(
                new ClassPathResource("keystore.jks"), "mySecretKey".toCharArray())
                .getKeyPair("coderider");
        converter.setKeyPair(keyPair);
        return converter;
    }

    @Override
    public void configure(AuthorizationServerEndpointsConfigurer endpoints) {
        endpoints.authenticationManager(this.authenticationManager)
                .accessTokenConverter(accessTokenConverter())
                .tokenGranter(OverrideTokenGranterConfigurer.tokenGranter(endpoints, this.authenticationManager))
                .addInterceptor(new HandlerInterceptorAdapter() {
                    @Override
                    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler,
                                ModelAndView modelAndView) {
                        if (modelAndView != null && modelAndView.getView() instanceof RedirectView) {
                            RedirectView redirect = (RedirectView) modelAndView.getView();
                            String url = redirect.getUrl();
                            if (url.contains("code=") || url.contains("error=")) {
                                HttpSession session = request.getSession(false);
                                if (session != null) {
                                    session.invalidate();
                                }
                            }
                        }
                    }
                })
        ;
    }

    @Bean
    public AuthorizationServerTokenServices tokenServices() {
        DefaultTokenServices service = new DefaultTokenServices();
        service.setClientDetailsService(clientDetailsService());
        service.setTokenStore(tokenStore());
        service.setTokenEnhancer(accessTokenConverter());
        service.setSupportRefreshToken(true);
        return service;
    }

    @Bean
    public AuthorizationCodeServices authorizationCodeServices(
                @Autowired @Qualifier("clientDetailsDataSource") DataSource clientDetailsDataSource) {
        return new JdbcAuthorizationCodeServices(clientDetailsDataSource);
    }

    @Override
    public void configure(AuthorizationServerSecurityConfigurer security) {
        security.tokenKeyAccess("permitAll()").checkTokenAccess("permitAll()")
                .allowFormAuthenticationForClients()
        ;
    }

    @Bean
    public ClientDetailsService clientDetailsService() {
        JdbcClientDetailsService clientDetailsService;
        clientDetailsService = new JdbcClientDetailsService(clientDetailsDataSource);
        clientDetailsService.setPasswordEncoder(passwordEncoder);
        return clientDetailsService;
    }

    @Override
    public void configure(ClientDetailsServiceConfigurer clients) throws Exception {
        clients.withClientDetails(clientDetailsService());
    }


}
